import torch
import torch.nn as nn
from model.encoder import Encoder
from model.decoder import Decoder
from model.attack import Attack
from noise_layers.gaussian_filter import GF
from noise_layers.gaussian_noise import Gaussian_Noise
from noise_layers.colorjitter import ColorJitter
from noise_layers.salt_pepper_noise import SP
from noise_layers.jpeg import Jpeg, JpegSS, JpegMask
from noise_layers.quantization import Quantization
"""
    空域图像隐写
"""

class EncoderDecoder(nn.Module):
    """
    Combines Encoder->Noiser->Decoder into single pipeline.
    The input is the cover image and the watermark message. The module inserts the watermark into the image
    (obtaining encoded_image), then applies Noise layers (obtaining noised_image), then passes the noised_image
    to the Decoder which tries to recover the watermark (called decoded_message). The module outputs
    a three-tuple: (encoded_image, noised_image, decoded_message)
    """
    def __init__(self, encoder_channels, decoder_channels, H, W, M, noise_channels=64, attack=True):
        super(EncoderDecoder, self).__init__()
        self.encoder = Encoder(encoder_channels, H, W, M) #channels=64
        self.noiser = Attack(noise_channels, attack)
        self.decoder = Decoder(decoder_channels, M) #channels=64
        # self.gaussian_filter = GF()
        # self.gaussian_noise = Gaussian_Noise()
        # self.colorJitter = ColorJitter('Brightness')
        # self.colorJitter = ColorJitter('Contrast')
        # self.colorJitter = ColorJitter('Saturation')
        # self.colorJitter = ColorJitter('Hue')
        # self.jp = JpegMask()
        # self.jp = Jpeg()
        # self.quantization = Quantization().to('cuda:3')
        # self.salt = SP()

    def forward(self, image, message):
        encoded_image = self.encoder(image, message)
        noise_image = self.noiser(encoded_image)
        # noise_image = self.colorJitter(encoded_image)
        # noise_image = self.gaussian_noise(encoded_image)
        # noise_image = self.salt(encoded_image)
        # noise_image = self.jp(encoded_image)
        # noise_image = self.quantization(encoded_image)
        # gf_image = self.gaussian_filter(noised_image)
        # noise_image = encoded_image
        decoded_message = self.decoder(noise_image)
        return encoded_image, noise_image, decoded_message